import numpy as np
# 生成系数矩阵A
import time
# w = np.array([1, 2, 3, 4, 5]) # x数组的权重值
# z = np.polyfit(x, y, 2, w=w) # 用2次多项式拟合
def lfit(x,y,m,w):
    if len(x)<=m:
        return False
    xishu = np.polyfit(x,y,m,w=w)
    p = np.poly1d(xishu) # 构造多项式
    yfit = p(x) # 拟合的y值
    yresid = y - yfit # 残差
    SSresid = sum(pow(yresid, 2)) # 残差平方和
    SStotal = len(y) * np.var(y) # 总体平均方差
    if SStotal==0:
        r2=1
    else:
        r2 = 1 - SSresid/SStotal # 拟合优度    
    # xishu=list(xishu)
    # xishu.reverse()
    return (xishu,r2,yfit)
def lfit0(x,y,m):
    m=m+1
    def gen_coefficient_matrix(X,m):
        N = len(X)
        A = []
        # 计算每一个方程的系数
        for i in range(m):
            a = []
        # 计算当前方程中的每一个系数
            for j in range(m):
                a.append(sum(X ** (i+j)))
            A.append(a)
        return A

    # 计算方程组的右端向量b
    def gen_right_vector(X, Y,m):
        N=len(X)
        b=[]
        for i in range(m):
            b.append(sum(X**i*y))
        return b

    A = gen_coefficient_matrix(x,m)
    b = gen_right_vector(x,y,m)
    xishu = np.linalg.solve(A, b)
    r=np.corrcoef(x,y)[0][1]
    return (xishu,r)
if __name__=="__main__":
    print(time.perf_counter())
    # _X = np.arange(0, 5, 0.1)
    # _Y = np.array([a0 + a1*x + a2*x**2 for x in _X])
    # y = 2+3x+4x^2
    x = np.arange(0,5,0.1)
    z = [2+3*x+4*x**2 for x in x]
    y = np.array([np.random.normal(z,3) for z in z])
    # plt.plot(x,y,'ro')
    # plt.show()
    print(y)
    print(lfit(x,y,2))
    # x = np.arange(0,5,0.1)
    # z = [3+5*x for x in x]
    # y = [np.random.normal(z,1) for z in z]
    # (xishu,r)=lfit(x,y,1)
    # print(xishu)
    # y_fitted=[xishu[0]+xishu[1]*x for x in x]
    # print(x)
    # print(y)
    # print(y_fitted)
    # print(time.perf_counter())
    # print(lfit(x,y,3))
    # r=np.corrcoef(x,y)
    # r0=r[0][1]
    # print(r,r0)
    # plt.plot(x, y, 'ro', _X, _Y, 'b', linewidth=2)
    # plt.title("y = {} + {}x + {}$x^2$ ".format(a0, a1, a2))
    # plt.show()